import socket
import serial
import argparse
import time
import sys

def udp_to_serial(udp_port, serial_port, baud_rate):
    """
    将UDP数据包转发到串口。

    参数:
    udp_port (int): 监听的UDP端口。
    serial_port (str): 串口端口。
    baud_rate (int): 串口波特率。
    """
    # 设置UDP服务器
    udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_sock.bind(("0.0.0.0", udp_port))

    # 设置串口
    ser = serial.Serial(serial_port, baud_rate)

    print(f"Listening for UDP packets on port {udp_port} and forwarding to serial port {serial_port} at {baud_rate} baud.")
    sys.stdout.flush()

    try:
        while True:
            # 接收UDP数据
            data, addr = udp_sock.recvfrom(1024)
            message = data.decode()
            print(f"Received data from {addr}: {message}")
            sys.stdout.flush()

            if message == "test":
                udp_sock.sendto(b"okay", addr)
                print(f"Sent 'okay' to {addr}")
                sys.stdout.flush()
            else:
                # 将数据写入串口
                ser.write(data)
                print(f"Forwarded data to serial port {serial_port}.")
                sys.stdout.flush()
            # time.sleep(0.05)
    except KeyboardInterrupt:
        print("Stopping server.")
        sys.stdout.flush()
    finally:
        udp_sock.close()
        ser.close()

if __name__ == "__main__":
    """
    主程序入口。解析命令行参数并启动UDP到串口的转发器。
    """
    parser = argparse.ArgumentParser(description="UDP to Serial Forwarder")
    parser.add_argument("--udp_port", type=int, required=True, help="UDP port to listen on")
    parser.add_argument("--serial_port", type=str, required=True, help="Serial port to write to")
    parser.add_argument("--baud_rate", type=int, required=True, help="Baud rate for the serial port")

    args = parser.parse_args()
    udp_to_serial(args.udp_port, args.serial_port, args.baud_rate)
